﻿

using NgpSharp;
using static TorchSharp.torch.nn;
using TorchSharp;
using Serilog;
using static TorchSharp.torch;

namespace NgpTrainConsole
{
    public class Program
    {
        public const string DatasetPath = "E:\\Projects\\NeRFDataset\\lego\\transforms_train.json";
        public const string EvalPath = "E:\\Projects\\NeRFDataset\\lego\\transforms_val.json";

        public static string OutputPath = Path.Combine(AppContext.BaseDirectory, "output");

        public static async Task Main(string[] args)
        {
            try
            {
                if(!Directory.Exists(OutputPath)) Directory.CreateDirectory(OutputPath);

                NeRFDataset dataset = await NeRFDataset.LoadFromTransformsJson(DatasetPath);
                dataset.Training = true;
                dataset.NumRays = 128;
                dataset.BatchSize = 1;
                var Logger = new Serilog.LoggerConfiguration()
                    .MinimumLevel.Debug()
                    .WriteTo.Console()
                    .CreateLogger();
                int iters = 20000;

                var criterion = HuberLoss(reduction: Reduction.None, delta:0.1).cuda();
                var model = new NeRFNetwork("ngp", density_thresh: 10).cuda();
                model.num_step = 512;
                var optimizer = torch.optim.Adam(model.GetParameters(), lr: 1e-2, beta1: 0.9, beta2: 0.99, eps: 1e-15);
                var scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, (iter) => Math.Pow(0.1, Math.Min((iter / (double)iters), 1)));
                var device = torch.device("cuda", 0);

                NeRFTrainer trainer = new NeRFTrainer("ngp", device, model, optimizer, 
                    criterion, scheduler, Logger);
                trainer.Workspace = OutputPath;

                int max_epoch = (int)Math.Ceiling(iters / (double)dataset.Count);

                trainer.Train(dataset.ToDataLoader(), max_epoch);

                NeRFDataset evalDataset = await NeRFDataset.LoadFromTransformsJson(EvalPath);
                evalDataset.Training = false;
                var min = evalDataset.BoxMin; var max = evalDataset.BoxMax;
                Tensor aabb = torch.from_array(new float[] { min.X, min.Y, min.Z, max.X, max.Y, max.Z });
                var loader = evalDataset.ToDataLoader(shuffle: true);
                foreach(var data in loader)
                {
                    trainer.EvalStep(data, aabb); //only ones
                    break;
                }

            }
            catch (Exception e)
            {
                Console.WriteLine(e);
                Console.ReadLine();
            }
            Console.ReadLine();
        }
    }
}
